#include "postprocess/object_track.h"

IdGenerator ObjectTrack::id_generator_{};
std::int32_t ObjectTrack::s_track_cached_history_size_maximum_ = 5;
float ObjectTrack::s_acceleration_noise_maximum_ = 4.0;
float ObjectTrack::s_speed_noise_maximum_ = 0.4;

ObjectTrack::ObjectTrack(TrackedObjectPtr obj) {
  init();
  Eigen::Vector3f initial_center_point = obj->center;
  Eigen::Vector3f initial_velocity = Eigen::Vector3f::Zero();
  filter_ = std::make_shared<KalmanFilter>();
  filter_->Initialize(initial_center_point, initial_velocity);

  belief_center_point_ = obj->center;
  belief_velocity_ = Eigen::Vector3f::Zero();
  const float uncertainty_factor = 5.0;
  belief_velocity_uncertainty_ =
      Eigen::Matrix3f::Identity() * uncertainty_factor;
  belief_acceleration_ = Eigen::Vector3f::Zero();
  belief_acceleration_uncertainty_ = Eigen::Matrix3f::Identity() * 2.0f;
  belief_state_ = Eigen::Vector4f::Zero();
  belief_state_covariance_ = Eigen::Matrix4f::Identity();

  current_object_ = obj;
}

ObjectTrack::~ObjectTrack() {}

void ObjectTrack::init() {
  idx_ = id_generator_.GetId();
  age_ = 1;
  total_visible_count_ = 1;
  consecutive_invisible_count_ = 0;
  period_ = 0.0;
}

Eigen::VectorXf ObjectTrack::Predict(double time_diff) {
  // Get the predict of filter
  // Eigen::VectorXf filter_predict = filter_->Predict(time_diff);
  // Get the predict of track
  Eigen::VectorXf track_predict;
  track_predict.resize(7);
  track_predict(0) = belief_center_point_(0) + belief_velocity_(0) * time_diff;
  track_predict(1) = belief_center_point_(1) + belief_velocity_(1) * time_diff;
  track_predict(2) = belief_center_point_(2) + belief_velocity_(2) * time_diff;
  track_predict(3) = belief_velocity_(0);
  track_predict(4) = belief_velocity_(1);
  track_predict(5) = belief_velocity_(2);
  track_predict(6) = time_diff;
  return track_predict;
}

void ObjectTrack::UpdateWithObject(TrackedObjectPtr* new_object,
                                   double time_diff) {
  // A. update object track
  // A.1 update filter
  filter_->UpdateWithObject((*new_object), current_object_, time_diff);
  filter_->GetState(&belief_center_point_, &belief_velocity_,
                    &belief_acceleration_);
  //  filter_->GetOnlineCovariance(&belief_velocity_uncertainty_);
  filter_->GetVelocityCovariance(&belief_velocity_uncertainty_);
  filter_->GetAccelerationCovariance(&belief_acceleration_uncertainty_);
  filter_->GetStateCovariance(&belief_state_covariance_);
  // NEED TO NOTICE: All the states would be collected mainly based on states
  // of tracked object. Thus, update tracked object when you update the state
  // of track !!!!!
  // A.2 update tracked object
  (*new_object)->center = belief_center_point_;
  (*new_object)->velocity = belief_velocity_;
  // (*new_object)->velocity_uncertainty = belief_velocity_uncertainty_;
  (*new_object)->acceleration = belief_acceleration_;
  // (*new_object)->acceleration_uncertainty = belief_acceleration_uncertainty_;
  // (*new_object)->state = belief_state_;
  // (*new_object)->state_covariance = belief_state_covariance_;

  // A.2 update track info
  ++age_;
  total_visible_count_++;
  consecutive_invisible_count_ = 0;
  period_ += time_diff;
  // A.3 update history
  // A.4 update history
  std::int32_t history_size = history_objects_.size();
  if (history_size >= s_track_cached_history_size_maximum_) {
    history_objects_.pop_front();
  }
  history_objects_.push_back(current_object_);
  current_object_ = *new_object;
  // B. smooth object track
  // B.1 smooth velocity
  // smooth_track_velocity((*new_object), time_diff);
  // B.2 smooth orientation
  //  SmoothTrackOrientation();
}

void ObjectTrack::UpdateWithoutObject(Eigen::VectorXf const& predict_state,
                                      double time_diff) {
  // A. Update object of track
  TrackedObjectPtr new_obj(new TrackedObject());
  // new_obj->clone(*current_object_);

  new_obj->id = current_object_->id;
  new_obj->center =
      Eigen::Vector3f{current_object_->center.x(), current_object_->center.y(),
                      current_object_->center.z()};
  new_obj->size =
      Eigen::Vector3f{current_object_->size.x(), current_object_->size.y(),
                      current_object_->size.z()};
  new_obj->heading = current_object_->heading;
  new_obj->r = current_object_->r;
  new_obj->score = current_object_->score;
  new_obj->label = current_object_->label;
  new_obj->polygon = current_object_->polygon;
  new_obj->velocity = Eigen::Vector3f{current_object_->velocity.x(),
                                      current_object_->velocity.y(),
                                      current_object_->velocity.z()};
  new_obj->acceleration = Eigen::Vector3f{current_object_->acceleration.x(),
                                          current_object_->acceleration.y(),
                                          current_object_->acceleration.z()};
  new_obj->association_score = current_object_->association_score;

  Eigen::Vector3f predicted_shift = predict_state.segment(3, 3) * time_diff;
  new_obj->center = current_object_->center + predicted_shift;

  new_obj->polygon.Clear();
  for (size_t i = 0; i < 4; ++i) {
    auto& point{new_obj->polygon.GetNext()};
    point.x += predicted_shift[0];
    point.x += predicted_shift[1];
    new_obj->polygon.Add();
  }

  // C. Update filter
  filter_->UpdateWithoutObject(time_diff);

  // D. Update states of track
  belief_center_point_ = new_obj->center;
  new_obj->velocity = belief_velocity_;
  // new_obj->velocity_uncertainty = belief_velocity_uncertainty_;
  new_obj->acceleration = belief_acceleration_;
  // new_obj->acceleration_uncertainty = belief_acceleration_uncertainty_;

  // E. update track info
  ++age_;
  consecutive_invisible_count_++;
  period_ += time_diff;
  // }

  // F. update history
  std::int32_t history_size = history_objects_.size();
  if (history_size >= s_track_cached_history_size_maximum_) {
    history_objects_.pop_front();
  }
  history_objects_.push_back(current_object_);
  current_object_ = new_obj;
}

/*****************Object_track_set*******************/
float ObjectTrackSet::s_track_visible_ratio_minimum_ = 0.6;
std::int32_t ObjectTrackSet::s_track_consecutive_invisible_maximum_ = 4;

ObjectTrackSet::ObjectTrackSet() {
  tracks_.reserve(300U);
  ObjectTrack::InitIdGenerator();
}

ObjectTrackSet::~ObjectTrackSet() { Clear(); }

bool ObjectTrackSet::SetTrackConsecutiveInvisibleMaximum(
    std::int32_t track_consecutive_invisible_maximum) {
  if (track_consecutive_invisible_maximum >= 0) {
    s_track_consecutive_invisible_maximum_ =
        track_consecutive_invisible_maximum;
    return true;
  }
  return false;
}

bool ObjectTrackSet::SetTrackVisibleRatioMinimum(
    float track_visible_ratio_minimum) {
  if (track_visible_ratio_minimum >= 0 && track_visible_ratio_minimum <= 1) {
    s_track_visible_ratio_minimum_ = track_visible_ratio_minimum;
    return true;
  }
  return false;
}

void ObjectTrackSet::Clear() { tracks_.clear(); }

std::int32_t ObjectTrackSet::RemoveLostTracks() {
  std::int32_t track_num = 0;
  for (std::uint16_t i = 0; i < tracks_.size(); ++i) {
    // A. remove tracks invisible ratio is less than given minimum
    if (tracks_[i]->age_ == 0) {
      continue;
    } else {
      float track_visible_ratio =
          tracks_[i]->total_visible_count_ * 1.0f / tracks_[i]->age_;
      if (track_visible_ratio < s_track_visible_ratio_minimum_) {
        continue;
      }
    }
    // B. remove tracks consecutive invisible count greater than given maximum
    std::int32_t track_consecutive_invisible_count =
        tracks_[i]->consecutive_invisible_count_;
    if (track_consecutive_invisible_count >
        s_track_consecutive_invisible_maximum_) {
      continue;
    }
    // C. change order, move those objects who needs to be removed to the end of
    // the vector
    if (i == track_num) {
      track_num++;
    } else {
      ObjectTrackPtr tmp{tracks_[i]};
      tracks_[i] = tracks_[track_num];
      tracks_[track_num] = tmp;
      track_num++;
    }
  }

  // Remove lost tracks
  std::int32_t no_removed = tracks_.size() - track_num;
  for (size_t i = track_num; i < tracks_.size(); ++i) {
    if (tracks_[i] != nullptr) {
      delete (tracks_[i]);
      tracks_[i] = nullptr;
    }
  }
  tracks_.resize(track_num);
  return no_removed;
}
